#include "SHM.h"

vtype SHM1_prob::func(const vtype& input, const vtype& args, dtype t)
{
    assert(input.size() == 2 && args.size() == 1); 
    vtype f(2); 
    f[theta] = input[omega]; 
    f[omega] = -input[theta] * args[mu] * args[mu]; 
    return f; 
}

vtype SHM2_prob::func(const vtype& input, const vtype& args, dtype t)
{
    assert(input.size() == 2 && args.size() == 2);
    vtype f(2);
    f[theta] = input[omega];
    f[omega] = -input[theta] * args[mu] * args[mu] - input[omega] * args[fei];
    return f;
}

vtype SHM3_prob::func(const vtype& input, const vtype& args, dtype t)
{
    assert(input.size() == 2 && args.size() == 4);
    vtype f(2);
    f[theta] = input[omega];
    f[omega] = -input[theta] * args[mu] * args[mu] - input[omega] * args[fei]
        + args[omega_d] * cos(args[omega_d] * t); 
    return f;
}

vtype NIPP::func(const vtype& input, const vtype& args, dtype t)
{
    assert(input.size() == 2 && args.size() == 1);
    vtype f(2);
    f[theta] = input[omega];
    f[omega] = -sin(input[theta]) * args[mu] * args[mu];
    return f; 
}


vtype euler_solver::step(dtype h, dtype t)
{
    vtype k = prob_->func(prob_->val(), prob_->arg(), 0);
    for (int i = 0; i < k.size(); ++i) {
        prob_->val(i) = prob_->val(i) + h * k[i]; 
    }
    return prob_->val(); 
}

vtype rk2_solver::step(dtype h, dtype t)
{
    vector<double> k1 = prob_->func(prob_->val(), prob_->arg(), t);
    vector<double> k2 = prob_->val(); 
    for (int i = 0; i < k2.size(); ++i) {
        k2[i] += h * k1[i];
    }
    k2 = prob_->func(k2, prob_->arg(), t);
    for (int i = 0; i < k2.size(); ++i) {
        prob_->val(i) = prob_->val(i) + h * (k1[i] + k2[i]) / 2;
    }
    return prob_->val();
}

vtype rk4_solver::step(dtype h, dtype t)
{
    vtype k1 = prob_->func(prob_->val(), prob_->arg(), t);
    vtype k2 = prob_->val();
    for (int i = 0; i < prob_->val().size(); ++i)
        k2[i] += k1[i] * h / 2;
    k2 = prob_->func(k2, prob_->arg(), t);
    vtype k3 = prob_->val();
    for (int i = 0; i < prob_->val().size(); ++i)
        k3[i] += k2[i] * h / 2;
    k3 = prob_->func(k3, prob_->arg(), t);
    vtype k4 = prob_->val();
    for (int i = 0; i < prob_->val().size(); ++i)
        k4[i] += k3[i] * h;
    k4 = prob_->func(k4, prob_->arg(), t);
    for (int i = 0; i < prob_->val().size(); ++i)
        prob_->val(i) = prob_->val(i) + h * (k1[i] + 2 * k2[i] + 2 * k3[i] + k4[i]) / 6;
    return prob_->val();
}

vtype DPP::func(const vtype& input, const vtype& args, dtype t)
{
    dtype
        a1 = (args[m1] + args[m2]) * args[l1],
        a2 = args[m2] * args[l2] * cos(input[theta1] - input[theta2]),
        a3 = args[m2] * args[l2] * sin(input[theta1] - input[theta2]),
        a4 = (args[m1] + args[m2]) * 9.8f * sin(input[theta1]); 
    dtype
        b1 = args[l1] * cos(input[theta1] - input[theta2]),
        b2 = args[l2],
        b3 = -args[l1] * sin(input[theta1] - input[theta2]),
        b4 = 9.8f * sin(input[theta2]); 
    vtype f(4); 
    f[omega1] = (
        a3 * b2 * input[omega2] * input[omega2] - a2 * b3 * input[omega1] * input[omega1]
        + a4 * b2 - a2 * b4
        ) / (a1 * b2 - a2 * b1); 
    f[omega2] = (
        a3 * b1 * input[omega2] * input[omega2] - a1 * b3 * input[omega1] * input[omega1]
        + a4 * b1 - a1 * b4
        ) / (a2 * b1 - a1 * b2); 
    f[theta1] = -input[omega1]; 
    f[theta2] = -input[omega2]; 
    return f; 
}
